{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c8734f831fdc4098",
   "metadata": {},
   "source": "# Tutorial for regression"
  },
  {
   "cell_type": "markdown",
   "id": "9c2b8cffaddbe46a",
   "metadata": {},
   "source": "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/scikit-learn-contrib/MAPIE/blob/master/notebooks/regression/exoplanets.ipynb)\n"
  },
  {
   "cell_type": "markdown",
   "id": "75824ecc30df3ee6",
   "metadata": {},
   "source": ""
  },
  {
   "cell_type": "markdown",
   "id": "ea2b2d2f054ca596",
   "metadata": {},
   "source": [
    "In this tutorial, we compare the prediction intervals estimated by MAPIE on a\n",
    "simple, one-dimensional, ground truth function\n",
    "\n",
    "$$\n",
    "f(x) = x \\sin(x)\n",
    "$$\n",
    "\n",
    "Throughout this tutorial, we will answer the following questions:\n",
    "\n",
    "- How well do the MAPIE strategies capture the aleatoric uncertainty existing in the data?\n",
    "\n",
    "- How do the prediction intervals estimated by the resampling strategies\n",
    "  evolve for new *out-of-distribution* data? \n",
    "\n",
    "- How do the prediction intervals vary between regressor models?\n",
    "\n",
    "Throughout this tutorial, we estimate the prediction intervals first using \n",
    "a polynomial function, and then using a boosting model, and a simple neural network. \n",
    "\n",
    "**For practical problems, we advise using the faster CV+ strategies. \n",
    "For conservative prediction interval estimates, you can alternatively \n",
    "use the CV-minmax strategies.**\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "735f0f18088f1691",
   "metadata": {},
   "source": "## 1. Estimating the aleatoric uncertainty of homoscedastic noisy data"
  },
  {
   "cell_type": "markdown",
   "id": "80b1bc51f26c2cab",
   "metadata": {},
   "source": [
    "Let's start by defining the $x \\times \\sin(x)$ function and another simple function\n",
    "that generates one-dimensional data with normal noise uniformely in a given interval."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "501fe8bf258d2250",
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "import numpy as np\n",
    "def x_sinx(x):\n",
    "    \"\"\"One-dimensional x*sin(x) function.\"\"\"\n",
    "    return x*np.sin(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "942facc9220b1292",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_1d_data_with_constant_noise(funct, min_x, max_x, n_samples, noise):\n",
    "    \"\"\"\n",
    "    Generate 1D noisy data uniformely from the given function \n",
    "    and standard deviation for the noise.\n",
    "    \"\"\"\n",
    "    np.random.seed(59)\n",
    "    X_train = np.linspace(min_x, max_x, n_samples)\n",
    "    np.random.shuffle(X_train)\n",
    "    X_test = np.linspace(min_x, max_x, n_samples*5)\n",
    "    y_train, y_mesh, y_test = funct(X_train), funct(X_test), funct(X_test)\n",
    "    y_train += np.random.normal(0, noise, y_train.shape[0])\n",
    "    y_test += np.random.normal(0, noise, y_test.shape[0])\n",
    "    return X_train.reshape(-1, 1), y_train, X_test.reshape(-1, 1), y_test, y_mesh"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c256110c91cc9",
   "metadata": {},
   "source": [
    "We first generate noisy one-dimensional data uniformely on an interval. \n",
    "Here, the noise is considered as *homoscedastic*, since it remains constant \n",
    "over $x$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92ec82ad9de7060b",
   "metadata": {},
   "outputs": [],
   "source": [
    "min_x, max_x, n_samples, noise = -5, 5, 600, 0.5\n",
    "X_train_conformalize, y_train_conformalize, X_test, y_test, y_mesh = get_1d_data_with_constant_noise(\n",
    "    x_sinx, min_x, max_x, n_samples, noise\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae8fa66a518466ab",
   "metadata": {},
   "source": "Let's visualize our noisy function. "
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a6071d2d0babe94",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "plt.xlabel(\"x\") ; plt.ylabel(\"y\")\n",
    "plt.scatter(X_train_conformalize, y_train_conformalize, color=\"C0\")\n",
    "_ = plt.plot(X_test, y_mesh, color=\"C1\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ddbcac859b82780",
   "metadata": {},
   "source": [
    "As mentioned previously, we fit our training data with a simple\n",
    "polynomial function. Here, we choose a degree equal to 10 so the function \n",
    "is able to perfectly fit $x \\times \\sin(x)$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42276094e2f8cfc4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.preprocessing import PolynomialFeatures\n",
    "from sklearn.linear_model import LinearRegression, QuantileRegressor\n",
    "from sklearn.pipeline import Pipeline\n",
    "\n",
    "degree_polyn = 10\n",
    "polyn_model = Pipeline(\n",
    "    [\n",
    "        (\"poly\", PolynomialFeatures(degree=degree_polyn)),\n",
    "        (\"linear\", LinearRegression())\n",
    "    ]\n",
    ")\n",
    "polyn_model_quant = Pipeline(\n",
    "    [\n",
    "        (\"poly\", PolynomialFeatures(degree=degree_polyn)),\n",
    "        (\"linear\", QuantileRegressor(\n",
    "                solver=\"highs\",\n",
    "                alpha=0,\n",
    "        ))\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27ce27e70a4efc5f",
   "metadata": {},
   "source": [
    "We then estimate the prediction intervals for all the strategies very easily with a\n",
    "`fit` and `predict` process. The prediction interval's lower and upper bounds\n",
    "are then saved in a DataFrame. Here, we set an alpha value of 0.05\n",
    "in order to obtain a 95% confidence for our prediction intervals."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30abe3c469c24fc5",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "from mapie.regression import CrossConformalRegressor, JackknifeAfterBootstrapRegressor, ConformalizedQuantileRegressor\n",
    "from sklearn.model_selection import train_test_split\n",
    "RANDOM_STATE = 1\n",
    "STRATEGIES = {\n",
    "    \"cv\": {\n",
    "        \"class\": CrossConformalRegressor,\n",
    "        \"init_params\": dict(method=\"base\", cv=10),\n",
    "    },\n",
    "    \"cv_plus\": {\n",
    "        \"class\": CrossConformalRegressor,\n",
    "        \"init_params\": dict(method=\"plus\", cv=10),\n",
    "    },\n",
    "    \"cv_minmax\": {\n",
    "        \"class\": CrossConformalRegressor,\n",
    "        \"init_params\": dict(method=\"minmax\", cv=10),\n",
    "    },\n",
    "    \"jackknife\": {\n",
    "        \"class\": CrossConformalRegressor,\n",
    "        \"init_params\": dict(method=\"base\", cv=-1),\n",
    "    },\n",
    "    \"jackknife_plus\": {\n",
    "        \"class\": CrossConformalRegressor,\n",
    "        \"init_params\": dict(method=\"plus\", cv=-1),\n",
    "    },\n",
    "    \"jackknife_minmax\": {\n",
    "        \"class\": CrossConformalRegressor,\n",
    "        \"init_params\": dict(method=\"minmax\", cv=-1),\n",
    "    },\n",
    "    \"jackknife_plus_ab\": {\n",
    "        \"class\": JackknifeAfterBootstrapRegressor,\n",
    "        \"init_params\": dict(method=\"plus\", resampling=50),\n",
    "    },\n",
    "    \"jackknife_minmax_ab\": {\n",
    "        \"class\": JackknifeAfterBootstrapRegressor,\n",
    "        \"init_params\": dict(method=\"minmax\", resampling=50),\n",
    "    },\n",
    "    \"conformalized_quantile_regression\": {\n",
    "        \"class\": ConformalizedQuantileRegressor,\n",
    "        \"init_params\": dict(),\n",
    "    },\n",
    "}\n",
    "y_pred, y_pis = {}, {}\n",
    "for strategy_name, strategy_params in STRATEGIES.items():\n",
    "    init_params = strategy_params[\"init_params\"]\n",
    "    class_ = strategy_params[\"class\"]\n",
    "    if strategy_name == \"conformalized_quantile_regression\":\n",
    "        X_train, X_conformalize, y_train, y_conformalize = (\n",
    "            train_test_split(\n",
    "                X_train_conformalize, y_train_conformalize,\n",
    "                test_size=0.3, random_state=RANDOM_STATE\n",
    "            )\n",
    "        )\n",
    "        mapie = class_(polyn_model_quant, confidence_level=0.95, **init_params)\n",
    "        mapie.fit(X_train, y_train)\n",
    "        mapie.conformalize(X_conformalize, y_conformalize)\n",
    "        y_pred[strategy_name], y_pis[strategy_name] = mapie.predict_interval(X_test)\n",
    "    else:\n",
    "        mapie = class_(\n",
    "            polyn_model, confidence_level=0.95, random_state=RANDOM_STATE, **init_params\n",
    "        )\n",
    "        mapie.fit_conformalize(X_train_conformalize, y_train_conformalize)\n",
    "        y_pred[strategy_name], y_pis[strategy_name] = mapie.predict_interval(X_test)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "caa9b4d3ff3795ad",
   "metadata": {},
   "source": [
    "Let’s now compare the target confidence intervals with the predicted intervals obtained \n",
    "with the Jackknife+, Jackknife-minmax, CV+, CV-minmax, Jackknife+-after-Boostrap, and conformalized quantile regression (CQR) strategies. Note that for the Jackknife-after-Bootstrap method, we call the :class:`mapie.subsample.Subsample` object that allows us to train bootstrapped models. Note also that the CQR method is called with :class:`MapieQuantileRegressor` with a \"split\" strategy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f58e17ee59c1a7ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_1d_data(\n",
    "    X_train,\n",
    "    y_train, \n",
    "    X_test,\n",
    "    y_test,\n",
    "    y_sigma,\n",
    "    y_pred, \n",
    "    y_pred_low, \n",
    "    y_pred_up,\n",
    "    ax=None,\n",
    "    title=None\n",
    "):\n",
    "    ax.set_xlabel(\"x\") ; ax.set_ylabel(\"y\")\n",
    "    ax.fill_between(X_test, y_pred_low, y_pred_up, alpha=0.3)\n",
    "    ax.scatter(X_train, y_train, color=\"red\", alpha=0.3, label=\"Training data\")\n",
    "    ax.plot(X_test, y_test, color=\"gray\", label=\"True confidence intervals\")\n",
    "    ax.plot(X_test, y_test - y_sigma, color=\"gray\", ls=\"--\")\n",
    "    ax.plot(X_test, y_test + y_sigma, color=\"gray\", ls=\"--\")\n",
    "    ax.plot(X_test, y_pred, color=\"blue\", alpha=0.5, label=\"Prediction intervals\")\n",
    "    if title is not None:\n",
    "        ax.set_title(title)\n",
    "    ax.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89a4b256fcd42de",
   "metadata": {},
   "outputs": [],
   "source": [
    "strategies = [\"jackknife_plus\", \"jackknife_minmax\", \"cv_plus\", \"cv_minmax\", \"jackknife_plus_ab\", \"conformalized_quantile_regression\"]\n",
    "n_figs = len(strategies)\n",
    "fig, axs = plt.subplots(3, 2, figsize=(9, 13))\n",
    "coords = [axs[0, 0], axs[0, 1], axs[1, 0], axs[1, 1], axs[2, 0], axs[2, 1]]\n",
    "for strategy, coord in zip(strategies, coords):\n",
    "    plot_1d_data(\n",
    "        X_train_conformalize.ravel(),\n",
    "        y_train_conformalize.ravel(),\n",
    "        X_test.ravel(),\n",
    "        y_mesh.ravel(),\n",
    "        np.full((X_test.shape[0]), 1.96*noise).ravel(),\n",
    "        y_pred[strategy].ravel(),\n",
    "        y_pis[strategy][:, 0, 0].ravel(),\n",
    "        y_pis[strategy][:, 1, 0].ravel(),\n",
    "        ax=coord,\n",
    "        title=strategy\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9dec04b3313b11b",
   "metadata": {},
   "source": [
    "At first glance, the four strategies give similar results and the\n",
    "prediction intervals are very close to the true confidence intervals.\n",
    "Let’s confirm this by comparing the prediction interval widths over\n",
    "$x$ between all strategies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d982e3d007828a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize=(7, 5))\n",
    "ax.axhline(1.96*2*noise, ls=\"--\", color=\"k\", label=\"True width\")\n",
    "for strategy in STRATEGIES:\n",
    "    ax.plot(X_test, y_pis[strategy][:, 1, 0] - y_pis[strategy][:, 0, 0], label=strategy)\n",
    "ax.set_xlabel(\"x\")\n",
    "ax.set_ylabel(\"Prediction Interval Width\")\n",
    "_ = ax.legend(fontsize=10, loc=[1, 0.4])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af2573b0b18e67ed",
   "metadata": {},
   "source": [
    "As expected, the prediction intervals estimated by the Naive method\n",
    "are slightly too narrow. The Jackknife, Jackknife+, CV, CV+, JaB, and J+aB give\n",
    "similar widths that are very close to the true width. On the other hand,\n",
    "the width estimated by Jackknife-minmax and CV-minmax are slightly too\n",
    "wide. Note that the widths given by the Naive, Jackknife, and CV strategies\n",
    "are constant because there is a single model used for prediction,\n",
    "perturbed models are ignored at prediction time.\n",
    "\n",
    "It's interesting to observe that CQR strategy offers more varying width,\n",
    "often giving much higher but also lower interval width than other methods, therefore,\n",
    "with homoscedastic noise, CQR would not be the preferred method."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1841c05f69895538",
   "metadata": {},
   "source": [
    "Let’s now compare the *effective* coverage, namely the fraction of test\n",
    "points whose true values lie within the prediction intervals, given by\n",
    "the different strategies. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e5539d2f6d8bcc9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "from mapie.metrics.regression import regression_coverage_score\n",
    "\n",
    "pd.DataFrame([\n",
    "    [\n",
    "        regression_coverage_score(\n",
    "            y_test, y_pis[strategy]\n",
    "        )[0],\n",
    "        (\n",
    "            y_pis[strategy][:, 1, 0] - y_pis[strategy][:, 0, 0]\n",
    "        ).mean()\n",
    "    ] for strategy in STRATEGIES\n",
    "], index=STRATEGIES, columns=[\"Coverage\", \"Width average\"]).round(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "477561e612a1539",
   "metadata": {},
   "source": [
    "All strategies except the Naive one give effective coverage close to the expected \n",
    "0.95 value (recall that alpha = 0.05), confirming the theoretical garantees."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8983519c3adc695d",
   "metadata": {},
   "source": "## 2. Estimating the aleatoric uncertainty of heteroscedastic noisy data"
  },
  {
   "cell_type": "markdown",
   "id": "9efe45da20158493",
   "metadata": {},
   "source": [
    "Let's define again the $x \\times \\sin(x)$ function and another simple function\n",
    "that generates one-dimensional data with normal noise uniformely in a given interval."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f77a8a8245be75f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def x_sinx(x):\n",
    "    \"\"\"One-dimensional x*sin(x) function.\"\"\"\n",
    "    return x*np.sin(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aaeaf039747e903d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_1d_data_with_heteroscedastic_noise(funct, min_x, max_x, n_samples, noise):\n",
    "    \"\"\"\n",
    "    Generate 1D noisy data uniformely from the given function \n",
    "    and standard deviation for the noise.\n",
    "    \"\"\"\n",
    "    np.random.seed(59)\n",
    "    X_train = np.linspace(min_x, max_x, n_samples)\n",
    "    np.random.shuffle(X_train)\n",
    "    X_test = np.linspace(min_x, max_x, n_samples*5)\n",
    "    y_train = funct(X_train) + (np.random.normal(0, noise, len(X_train)) * X_train)\n",
    "    y_test = funct(X_test) + (np.random.normal(0, noise, len(X_test)) * X_test)\n",
    "    y_mesh = funct(X_test)\n",
    "    return X_train.reshape(-1, 1), y_train, X_test.reshape(-1, 1), y_test, y_mesh"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f30b4873b4c389d7",
   "metadata": {},
   "source": [
    "We first generate noisy one-dimensional data uniformely on an interval. \n",
    "Here, the noise is considered as *heteroscedastic*, since it will increase linearly with $x$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10fe19be966bb439",
   "metadata": {},
   "outputs": [],
   "source": [
    "min_x, max_x, n_samples, noise = 0, 5, 300, 0.5\n",
    "X_train_conformalize, y_train_conformalize, X_test, y_test, y_mesh = get_1d_data_with_heteroscedastic_noise(\n",
    "    x_sinx, min_x, max_x, n_samples, noise\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4948e46591058b33",
   "metadata": {},
   "source": "Let's visualize our noisy function. As x increases, the data becomes more noisy."
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84e529e5f9532408",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "plt.xlabel(\"x\") ; plt.ylabel(\"y\")\n",
    "plt.scatter(X_train_conformalize, y_train_conformalize, color=\"C0\")\n",
    "_ = plt.plot(X_test, y_mesh, color=\"C1\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59a3bcbbdf6222ea",
   "metadata": {},
   "source": [
    "As mentioned previously, we fit our training data with a simple\n",
    "polynomial function. Here, we choose a degree equal to 10 so the function \n",
    "is able to perfectly fit $x \\times \\sin(x)$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "656266998725ec6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.preprocessing import PolynomialFeatures\n",
    "from sklearn.linear_model import LinearRegression, QuantileRegressor\n",
    "from sklearn.pipeline import Pipeline\n",
    "\n",
    "degree_polyn = 10\n",
    "polyn_model = Pipeline(\n",
    "    [\n",
    "        (\"poly\", PolynomialFeatures(degree=degree_polyn)),\n",
    "        (\"linear\", LinearRegression())\n",
    "    ]\n",
    ")\n",
    "polyn_model_quant = Pipeline(\n",
    "    [\n",
    "        (\"poly\", PolynomialFeatures(degree=degree_polyn)),\n",
    "        (\"linear\", QuantileRegressor(\n",
    "                solver=\"highs\",\n",
    "                alpha=0,\n",
    "        ))\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b272d6f24a7c0ec",
   "metadata": {},
   "source": [
    "We then estimate the prediction intervals for all the strategies very easily with a\n",
    "`fit` and `predict` process. The prediction interval's lower and upper bounds\n",
    "are then saved in a DataFrame. Here, we set an alpha value of 0.05\n",
    "in order to obtain a 95% confidence for our prediction intervals."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1017289aa7f2b84",
   "metadata": {},
   "outputs": [],
   "source": [
    "from mapie.regression import (\n",
    "    CrossConformalRegressor,\n",
    "    JackknifeAfterBootstrapRegressor,\n",
    "    ConformalizedQuantileRegressor\n",
    ")\n",
    "from sklearn.model_selection import train_test_split\n",
    "RANDOM_STATE = 1\n",
    "STRATEGIES = {\n",
    "    \"cv\": {\n",
    "        \"class\": CrossConformalRegressor,\n",
    "        \"init_params\": dict(method=\"base\", cv=10),\n",
    "    },\n",
    "    \"cv_plus\": {\n",
    "        \"class\": CrossConformalRegressor,\n",
    "        \"init_params\": dict(method=\"plus\", cv=10),\n",
    "    },\n",
    "    \"cv_minmax\": {\n",
    "        \"class\": CrossConformalRegressor,\n",
    "        \"init_params\": dict(method=\"minmax\", cv=10),\n",
    "    },\n",
    "    \"jackknife\": {\n",
    "        \"class\": CrossConformalRegressor,\n",
    "        \"init_params\": dict(method=\"base\", cv=-1),\n",
    "    },\n",
    "    \"jackknife_plus\": {\n",
    "        \"class\": CrossConformalRegressor,\n",
    "        \"init_params\": dict(method=\"plus\", cv=-1),\n",
    "    },\n",
    "    \"jackknife_minmax\": {\n",
    "        \"class\": CrossConformalRegressor,\n",
    "        \"init_params\": dict(method=\"minmax\", cv=-1),\n",
    "    },\n",
    "    \"jackknife_plus_ab\": {\n",
    "        \"class\": JackknifeAfterBootstrapRegressor,\n",
    "        \"init_params\": dict(method=\"plus\", resampling=50),\n",
    "    },\n",
    "    \"jackknife_minmax_ab\": {\n",
    "        \"class\": JackknifeAfterBootstrapRegressor,\n",
    "        \"init_params\": dict(method=\"minmax\", resampling=50),\n",
    "    },\n",
    "    \"conformalized_quantile_regression\": {\n",
    "        \"class\": ConformalizedQuantileRegressor,\n",
    "        \"init_params\": dict(),\n",
    "    },\n",
    "}\n",
    "y_pred, y_pis = {}, {}\n",
    "for strategy_name, strategy_params in STRATEGIES.items():\n",
    "    init_params = strategy_params[\"init_params\"]\n",
    "    class_ = strategy_params[\"class\"]\n",
    "    if strategy_name == \"conformalized_quantile_regression\":\n",
    "        X_train, X_conformalize, y_train, y_conformalize = (\n",
    "            train_test_split(\n",
    "                X_train_conformalize, y_train_conformalize,\n",
    "                test_size=0.3, random_state=RANDOM_STATE\n",
    "            )\n",
    "        )\n",
    "        mapie = class_(polyn_model_quant, confidence_level=0.95, **init_params)\n",
    "        mapie.fit(X_train, y_train)\n",
    "        mapie.conformalize(X_conformalize, y_conformalize)\n",
    "        y_pred[strategy_name], y_pis[strategy_name] = mapie.predict_interval(X_test)\n",
    "    else:\n",
    "        mapie = class_(\n",
    "            polyn_model, confidence_level=0.95, random_state=RANDOM_STATE, **init_params\n",
    "        )\n",
    "        mapie.fit_conformalize(X_train_conformalize, y_train_conformalize)\n",
    "        y_pred[strategy_name], y_pis[strategy_name] = mapie.predict_interval(X_test)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fced1097252847fa",
   "metadata": {},
   "source": "Once again, let’s compare the target confidence intervals with prediction intervals obtained with the Jackknife+, Jackknife-minmax, CV+, CV-minmax, Jackknife+-after-Boostrap, and CQR strategies."
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96b8cf37c29d2578",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_1d_data(\n",
    "    X_train,\n",
    "    y_train, \n",
    "    X_test,\n",
    "    y_test,\n",
    "    y_sigma,\n",
    "    y_pred, \n",
    "    y_pred_low, \n",
    "    y_pred_up,\n",
    "    ax=None,\n",
    "    title=None\n",
    "):\n",
    "    ax.set_xlabel(\"x\") ; ax.set_ylabel(\"y\")\n",
    "    ax.fill_between(X_test, y_pred_low, y_pred_up, alpha=0.3)\n",
    "    ax.scatter(X_train, y_train, color=\"red\", alpha=0.3, label=\"Training data\")\n",
    "    ax.plot(X_test, y_test, color=\"gray\", label=\"True confidence intervals\")\n",
    "    ax.plot(X_test, y_test - y_sigma, color=\"gray\", ls=\"--\")\n",
    "    ax.plot(X_test, y_test + y_sigma, color=\"gray\", ls=\"--\")\n",
    "    ax.plot(X_test, y_pred, color=\"blue\", alpha=0.5, label=\"Prediction intervals\")\n",
    "    if title is not None:\n",
    "        ax.set_title(title)\n",
    "    ax.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "135b2a752c64f13",
   "metadata": {},
   "outputs": [],
   "source": [
    "strategies = [\"jackknife_plus\", \"jackknife_minmax\", \"cv_plus\", \"cv_minmax\", \"jackknife_plus_ab\", \"conformalized_quantile_regression\"]\n",
    "n_figs = len(strategies)\n",
    "fig, axs = plt.subplots(3, 2, figsize=(9, 13))\n",
    "coords = [axs[0, 0], axs[0, 1], axs[1, 0], axs[1, 1], axs[2, 0], axs[2, 1]]\n",
    "for strategy, coord in zip(strategies, coords):\n",
    "    plot_1d_data(\n",
    "        X_train_conformalize.ravel(),\n",
    "        y_train_conformalize.ravel(),\n",
    "        X_test.ravel(),\n",
    "        y_mesh.ravel(),\n",
    "        (1.96*noise*X_test).ravel(),\n",
    "        y_pred[strategy].ravel(),\n",
    "        y_pis[strategy][:, 0, 0].ravel(),\n",
    "        y_pis[strategy][:, 1, 0].ravel(),\n",
    "        ax=coord,\n",
    "        title=strategy\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3cdd7c00a1299ce4",
   "metadata": {},
   "source": [
    "We can observe that all of the strategies except CQR seem to have similar constant prediction intervals. \n",
    "On the other hand, the CQR strategy offers a solution that adapts the prediction\n",
    "intervals to the local noise."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "370fc0f13719e397",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize=(7, 5))\n",
    "ax.plot(X_test, 1.96*2*noise*X_test, ls=\"--\", color=\"k\", label=\"True width\")\n",
    "for strategy in STRATEGIES:\n",
    "    ax.plot(X_test, y_pis[strategy][:, 1, 0] - y_pis[strategy][:, 0, 0], label=strategy)\n",
    "ax.set_xlabel(\"x\")\n",
    "ax.set_ylabel(\"Prediction Interval Width\")\n",
    "_ = ax.legend(fontsize=10, loc=[1, 0.4])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4ac9520e499aeeb6",
   "metadata": {},
   "source": "One can observe that all the strategies behave in a similar way as in the first example shown previously. One exception is the CQR method which takes into account the heteroscedasticity of the data. In this method we observe very low interval widths at low values of $x$. This is the only method that even slightly follows the true width, and therefore is the preferred method for heteroscedastic data. Notice also that the true width is greater (lower) than the predicted width from the other methods at $x \\gtrapprox 3$ ($x \\leq 3$). This means that while the marginal coverage correct for these methods, the conditional coverage is likely not guaranteed as we will observe in the next figure."
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aee32fe87eae8f37",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_heteroscedastic_coverage(y_test, y_pis, STRATEGIES, bins):\n",
    "    recap ={}\n",
    "    for i in range(len(bins)-1):\n",
    "        bin1, bin2 = bins[i], bins[i+1]\n",
    "        name = f\"[{bin1}, {bin2}]\"\n",
    "        recap[name] = []\n",
    "        for strategy in STRATEGIES:\n",
    "            indices = np.where((X_test>=bins[i])*(X_test<=bins[i+1]))\n",
    "            y_test_trunc = np.take(y_test, indices)\n",
    "            y_low_ = np.take(y_pis[strategy][:, 0, 0], indices)\n",
    "            y_high_ = np.take(y_pis[strategy][:, 1, 0], indices)\n",
    "            score_coverage = regression_coverage_score(y_test_trunc[0], np.column_stack((y_low_[0], y_high_[0])))[0]\n",
    "            recap[name].append(score_coverage)\n",
    "    recap_df = pd.DataFrame(recap, index=STRATEGIES)\n",
    "    return recap_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a98463c53a5affb1",
   "metadata": {},
   "outputs": [],
   "source": [
    "bins = [0, 1, 2, 3, 4, 5]\n",
    "heteroscedastic_coverage = get_heteroscedastic_coverage(y_test, y_pis, STRATEGIES, bins)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3a30acece593ebb",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "heteroscedastic_coverage.T.plot.bar(figsize=(12, 4), alpha=0.7)\n",
    "plt.axhline(0.95, ls=\"--\", color=\"k\")\n",
    "plt.ylabel(\"Conditional coverage\")\n",
    "plt.xlabel(\"x bins\")\n",
    "plt.xticks(rotation=0)\n",
    "plt.ylim(0.8, 1.0)\n",
    "plt.legend(loc=[1, 0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "67f82b27c6d64a9",
   "metadata": {},
   "source": [
    "Let’s now conclude by summarizing the *effective* coverage, namely the fraction of test\n",
    "points whose true values lie within the prediction intervals, given by\n",
    "the different strategies. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3fdafd157a4ea603",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "pd.DataFrame([\n",
    "    [\n",
    "        regression_coverage_score(\n",
    "            y_test, y_pis[strategy]\n",
    "        )[0],\n",
    "        (\n",
    "            y_pis[strategy][:, 1, 0] - y_pis[strategy][:, 0, 0]\n",
    "        ).mean()\n",
    "    ] for strategy in STRATEGIES\n",
    "], index=STRATEGIES, columns=[\"Coverage\", \"Width average\"]).round(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8989ce16ddc75cf1",
   "metadata": {},
   "source": "All the strategies have the wanted coverage, however, we notice that the CQR strategy has much lower interval width than all the other methods, therefore, with heteroscedastic noise, CQR would be the preferred method."
  },
  {
   "cell_type": "markdown",
   "id": "2be2b247c689f365",
   "metadata": {},
   "source": "## 3. Estimating the epistemic uncertainty of out-of-distribution data"
  },
  {
   "cell_type": "markdown",
   "id": "3b0877e878acf2ba",
   "metadata": {},
   "source": [
    "Let’s now consider one-dimensional data without noise, but normally distributed.\n",
    "The goal is to explore how the prediction intervals evolve for new data \n",
    "that lie outside the distribution of the training data in order to see how the strategies\n",
    "can capture the *epistemic* uncertainty. \n",
    "For a comparison of the epistemic and aleatoric uncertainties, please have a look at this\n",
    "[source](https://en.wikipedia.org/wiki/Uncertainty_quantification)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "160f8e56f8867e09",
   "metadata": {},
   "source": "Lets\" start by generating and showing the data. "
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed8093dbcdf1288b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_1d_data_with_normal_distrib(funct, mu, sigma, n_samples, noise):\n",
    "    \"\"\"\n",
    "    Generate noisy 1D data with normal distribution from given function \n",
    "    and noise standard deviation.\n",
    "    \"\"\"\n",
    "    np.random.seed(59)\n",
    "    X_train = np.random.normal(mu, sigma, n_samples)\n",
    "    X_test = np.arange(mu-4*sigma, mu+4*sigma, sigma/20.)\n",
    "    y_train, y_mesh, y_test = funct(X_train), funct(X_test), funct(X_test)\n",
    "    y_train += np.random.normal(0, noise, y_train.shape[0])\n",
    "    y_test += np.random.normal(0, noise, y_test.shape[0])\n",
    "    return X_train.reshape(-1, 1), y_train, X_test.reshape(-1, 1), y_test, y_mesh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f3b168724270fff",
   "metadata": {},
   "outputs": [],
   "source": [
    "mu = 0 ; sigma = 2 ; n_samples = 1000 ; noise = 0.\n",
    "X_train_conformalize, y_train_conformalize, X_test, y_test, y_mesh = get_1d_data_with_normal_distrib(\n",
    "    x_sinx, mu, sigma, n_samples, noise\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eeee5e3efa201fba",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.xlabel(\"x\") ; plt.ylabel(\"y\")\n",
    "plt.scatter(X_train_conformalize, y_train_conformalize, color=\"C0\")\n",
    "_ = plt.plot(X_test, y_test, color=\"C1\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8348b36767f96f5",
   "metadata": {},
   "source": [
    "As before, we estimate the prediction intervals using a polynomial\n",
    "function of degree 10 and show the results for the Jackknife+ and CV+\n",
    "strategies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "996742c4c6c4998a",
   "metadata": {},
   "outputs": [],
   "source": [
    "polyn_model_quant = Pipeline(\n",
    "    [\n",
    "        (\"poly\", PolynomialFeatures(degree=degree_polyn)),\n",
    "        (\"linear\", QuantileRegressor(\n",
    "                solver=\"highs-ds\",\n",
    "                alpha=0,\n",
    "        ))\n",
    "    ]\n",
    ")\n",
    "from mapie.regression import (\n",
    "    CrossConformalRegressor,\n",
    "    JackknifeAfterBootstrapRegressor,\n",
    "    ConformalizedQuantileRegressor\n",
    ")\n",
    "from sklearn.model_selection import train_test_split\n",
    "RANDOM_STATE = 1\n",
    "STRATEGIES = {\n",
    "    \"cv\": {\n",
    "        \"class\": CrossConformalRegressor,\n",
    "        \"init_params\": dict(method=\"base\", cv=10),\n",
    "    },\n",
    "    \"cv_plus\": {\n",
    "        \"class\": CrossConformalRegressor,\n",
    "        \"init_params\": dict(method=\"plus\", cv=10),\n",
    "    },\n",
    "    \"cv_minmax\": {\n",
    "        \"class\": CrossConformalRegressor,\n",
    "        \"init_params\": dict(method=\"minmax\", cv=10),\n",
    "    },\n",
    "    \"jackknife\": {\n",
    "        \"class\": CrossConformalRegressor,\n",
    "        \"init_params\": dict(method=\"base\", cv=-1),\n",
    "    },\n",
    "    \"jackknife_plus\": {\n",
    "        \"class\": CrossConformalRegressor,\n",
    "        \"init_params\": dict(method=\"plus\", cv=-1),\n",
    "    },\n",
    "    \"jackknife_minmax\": {\n",
    "        \"class\": CrossConformalRegressor,\n",
    "        \"init_params\": dict(method=\"minmax\", cv=-1),\n",
    "    },\n",
    "    \"jackknife_plus_ab\": {\n",
    "        \"class\": JackknifeAfterBootstrapRegressor,\n",
    "        \"init_params\": dict(method=\"plus\", resampling=50),\n",
    "    },\n",
    "    \"jackknife_minmax_ab\": {\n",
    "        \"class\": JackknifeAfterBootstrapRegressor,\n",
    "        \"init_params\": dict(method=\"minmax\", resampling=50),\n",
    "    },\n",
    "    \"conformalized_quantile_regression\": {\n",
    "        \"class\": ConformalizedQuantileRegressor,\n",
    "        \"init_params\": dict(),\n",
    "    },\n",
    "}\n",
    "y_pred, y_pis = {}, {}\n",
    "for strategy_name, strategy_params in STRATEGIES.items():\n",
    "    init_params = strategy_params[\"init_params\"]\n",
    "    class_ = strategy_params[\"class\"]\n",
    "    if strategy_name == \"conformalized_quantile_regression\":\n",
    "        X_train, X_conformalize, y_train, y_conformalize = (\n",
    "            train_test_split(\n",
    "                X_train_conformalize, y_train_conformalize,\n",
    "                test_size=0.3, random_state=RANDOM_STATE\n",
    "            )\n",
    "        )\n",
    "        mapie = class_(polyn_model_quant, confidence_level=0.95, **init_params)\n",
    "        mapie.fit(X_train, y_train)\n",
    "        mapie.conformalize(X_conformalize, y_conformalize)\n",
    "        y_pred[strategy_name], y_pis[strategy_name] = mapie.predict_interval(X_test)\n",
    "    else:\n",
    "        mapie = class_(\n",
    "            polyn_model, confidence_level=0.95, random_state=RANDOM_STATE, **init_params\n",
    "        )\n",
    "        mapie.fit_conformalize(X_train_conformalize, y_train_conformalize)\n",
    "        y_pred[strategy_name], y_pis[strategy_name] = mapie.predict_interval(X_test)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1543e27580776b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "strategies = [\"jackknife_plus\", \"jackknife_minmax\", \"cv_plus\", \"cv_minmax\", \"jackknife_plus_ab\", \"conformalized_quantile_regression\"]\n",
    "n_figs = len(strategies)\n",
    "fig, axs = plt.subplots(3, 2, figsize=(9, 13))\n",
    "coords = [axs[0, 0], axs[0, 1], axs[1, 0], axs[1, 1], axs[2, 0], axs[2, 1]]\n",
    "for strategy, coord in zip(strategies, coords): \n",
    "    plot_1d_data(\n",
    "        X_train_conformalize.ravel(),\n",
    "        y_train_conformalize.ravel(), \n",
    "        X_test.ravel(),\n",
    "        y_mesh.ravel(),\n",
    "        1.96*noise, \n",
    "        y_pred[strategy].ravel(),\n",
    "        y_pis[strategy][:, 0, :].ravel(),\n",
    "        y_pis[strategy][:, 1, :].ravel(), \n",
    "        ax=coord,\n",
    "        title=strategy\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dfa3dba205bf1552",
   "metadata": {},
   "source": [
    "At first glance, our polynomial function does not give accurate\n",
    "predictions with respect to the true function when $|x > 6|$. \n",
    "The prediction intervals estimated with the Jackknife+ do not seem to \n",
    "increase significantly, unlike the CV+ method whose prediction intervals\n",
    "capture a high uncertainty when $x > 6$."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6500e2211b39eaa3",
   "metadata": {},
   "source": "Let's now compare the prediction interval widths between all strategies. \n"
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a54835578c5fa9d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize=(7, 5))\n",
    "ax.set_yscale(\"log\")\n",
    "for strategy in STRATEGIES:\n",
    "    ax.plot(X_test, y_pis[strategy][:, 1, 0] - y_pis[strategy][:, 0, 0], label=strategy)\n",
    "ax.set_xlabel(\"x\")\n",
    "ax.set_ylabel(\"Prediction Interval Width\")\n",
    "ax.legend(fontsize=10, loc=[1, 0.4]);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d67879f95702c5b1",
   "metadata": {},
   "source": [
    "The prediction interval widths start to increase exponentially\n",
    "for $|x| > 4$ for the CV+, CV-minmax, Jackknife-minmax, and quantile\n",
    "strategies. On the other hand, the prediction intervals estimated by\n",
    "Jackknife+ remain roughly constant until $|x| \\sim 5$ before\n",
    "increasing.\n",
    "The CQR strategy seems to perform well, however, on the extreme values\n",
    "of the data the quantile regression fails to give reliable results as it outputs\n",
    "negative value for the prediction intervals. This occurs because the quantile \n",
    "regressor with quantile $1 - \\alpha/2$ gives higher values than the quantile \n",
    "regressor with quantile $\\alpha/2$. Note that a warning will be issued when\n",
    "this occurs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9f91a11a3ec4c3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "pd.DataFrame([\n",
    "    [\n",
    "        regression_coverage_score(\n",
    "            y_test, y_pis[strategy]\n",
    "        )[0],\n",
    "        (\n",
    "            y_pis[strategy][:, 1, 0] - y_pis[strategy][:, 0, 0]\n",
    "        ).mean()\n",
    "    ] for strategy in STRATEGIES\n",
    "], index=STRATEGIES, columns=[\"Coverage\", \"Width average\"]).round(3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f09ee2afcecccd7",
   "metadata": {},
   "source": [
    "In conclusion, the Jackknife-minmax, CV+, CV-minmax, or Jackknife-minmax-ab strategies are more\n",
    "conservative than the Jackknife+ strategy, and tend to result in more\n",
    "reliable coverages for *out-of-distribution* data. It is therefore\n",
    "advised to use the three former strategies for predictions with new\n",
    "out-of-distribution data.\n",
    "Note however that there are no theoretical guarantees on the coverage level \n",
    "for out-of-distribution data.\n",
    "Here it's important to note that the CQR strategy should not be taken into account for\n",
    "width prediction, and it is abundantly clear from the negative width coverage that\n",
    "is observed in these results."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98f2f0712bcf8809",
   "metadata": {},
   "source": "## 4. Estimating the uncertainty with different sklearn-compatible regressors"
  },
  {
   "cell_type": "markdown",
   "id": "9db3ad0c98c96ce0",
   "metadata": {},
   "source": [
    "MAPIE can be used with any kind of sklearn-compatible regressor. Here, we\n",
    "illustrate this by comparing the prediction intervals estimated by the CV+ method using\n",
    "different models:\n",
    "\n",
    "- the same polynomial function as before.\n",
    " \n",
    "- a XGBoost model using the Scikit-learn API.\n",
    "\n",
    "- a simple neural network, a Multilayer Perceptron with three dense layers, using the KerasRegressor wrapper.\n",
    "\n",
    "Once again, let’s use our noisy one-dimensional data obtained from a\n",
    "uniform distribution."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b5714826953980b",
   "metadata": {},
   "outputs": [],
   "source": [
    "min_x, max_x, n_samples, noise = -5, 5, 100, 0.5\n",
    "X_train_conformalize, y_train_conformalize, X_test, y_test, y_mesh = get_1d_data_with_constant_noise(\n",
    "    x_sinx, min_x, max_x, n_samples, noise\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d195b3bac5c259c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.xlabel(\"x\") ; plt.ylabel(\"y\")\n",
    "plt.plot(X_test, y_mesh, color=\"C1\")\n",
    "_ = plt.scatter(X_train_conformalize, y_train_conformalize)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16865b902183c931",
   "metadata": {},
   "source": [
    "Let's then define the models. The boosing model considers 100 shallow trees with a max depth of 2 while\n",
    "the Multilayer Perceptron has two hidden dense layers with 20 neurons each followed by a relu activation.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb0eaa563bc3fb77",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.neural_network import MLPRegressor\n",
    "\n",
    "mlp_model = MLPRegressor(activation=\"relu\", random_state=RANDOM_STATE, max_iter=500)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec7edf534851d2e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "polyn_model = Pipeline(\n",
    "    [\n",
    "        (\"poly\", PolynomialFeatures(degree=degree_polyn)),\n",
    "        (\"linear\", LinearRegression())\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7a93953164a3cb1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.ensemble import GradientBoostingRegressor\n",
    "xgb_model = GradientBoostingRegressor(\n",
    "    max_depth=2,\n",
    "    n_estimators=100,\n",
    "    random_state=RANDOM_STATE,\n",
    "    learning_rate=0.1,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "240e271869a10197",
   "metadata": {},
   "source": [
    "Let's now use MAPIE to estimate the prediction intervals using the CV+ method \n",
    "and compare their prediction interval."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f5c597d68f7d2b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "models = [polyn_model, xgb_model, mlp_model]\n",
    "model_names = [\"polyn\", \"xgb\", \"mlp\"]\n",
    "prediction_interval = {}\n",
    "for name, model in zip(model_names, models):\n",
    "    mapie = CrossConformalRegressor(model, method=\"plus\", cv=5, confidence_level=0.95)\n",
    "    mapie.fit_conformalize(X_train_conformalize, y_train_conformalize)\n",
    "    y_pred[name], y_pis[name] = mapie.predict_interval(X_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f9ada3b67cc4aea",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(1, 3, figsize=(20, 6))\n",
    "for name, ax in zip(model_names, axs):\n",
    "    plot_1d_data(\n",
    "        X_train_conformalize.ravel(),\n",
    "        y_train_conformalize.ravel(),\n",
    "        X_test.ravel(),\n",
    "        y_mesh.ravel(),\n",
    "        1.96*noise,\n",
    "        y_pred[name].ravel(),\n",
    "        y_pis[name][:, 0, 0].ravel(),\n",
    "        y_pis[name][:, 1, 0].ravel(),\n",
    "        ax=ax,\n",
    "        title=name\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3000e29aac12022",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize=(7, 5))\n",
    "for name in model_names:\n",
    "    ax.plot(X_test, y_pis[name][:, 1, 0] - y_pis[name][:, 0, 0])\n",
    "ax.axhline(1.96*2*noise, ls=\"--\", color=\"k\")\n",
    "ax.set_xlabel(\"x\")\n",
    "ax.set_ylabel(\"Prediction Interval Width\")\n",
    "ax.legend(model_names + [\"True width\"], fontsize=8);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4514a0f641528686",
   "metadata": {},
   "source": [
    "As expected with the CV+ method, the prediction intervals are a bit \n",
    "conservative since they are slightly wider than the true intervals.\n",
    "However, the CV+ method on the three models gives very promising results \n",
    "since the prediction intervals closely follow the true intervals with $x$."
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "formats": "ipynb,md"
  },
  "kernelspec": {
   "display_name": ".venv-doc",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
